//
//  NSTreeController+WCExtensions.m
//  WabbitCode
//
//  Created by William Towe on 1/28/09.
//  Copyright 2009 Revolution Software. All rights reserved.
//

#import "NSTreeController+WCExtensions.h"
#import "NSArray+WCExtensions.h"
#import "NSTreeNode+WCExtensions.h"
#import "AIArrayAdditions.h"
#import "WCProject.h"


@implementation NSTreeController (WCExtensions)
- (NSTreeNode *)selectedNode; {
	NSArray *nodes = [self selectedNodes];
	return ([nodes count])?[nodes firstObject]:nil;
}
- (id)selectedRepresentedObject; {
	return [[self selectedNode] representedObject];
}
- (NSArray *)selectedRepresentedObjects; {
	return [[self selectedNodes] valueForKey:@"representedObject"];
}

- (NSTreeNode *)treeNodeForRepresentedObject:(id)object; {
	for (NSTreeNode *node in [self flattenedNodes]) {
		if ([[node representedObject] isEqualTo:object])
			return node;
	}
	return nil;
}
- (NSArray *)treeNodesForRepresentedObjects:(NSArray *)objects; {
	NSMutableArray *nodes = [NSMutableArray array];
	NSArray *flattenedNodes = [self flattenedNodes];
	
	for (id object in objects) {
		for (NSTreeNode *node in flattenedNodes) {
			if ([[node representedObject] isEqualTo:object]) {
				[nodes addObject:node];
				break;
			}
		}
	}
	return [[nodes copy] autorelease];
}

- (NSArray *)flattenedNodes; {
	NSMutableArray *nodes = [NSMutableArray array];
	for (NSTreeNode *node in [[self arrangedObjects] childNodes]) {
		[nodes addObject:node];
		if (![[node valueForKey:[self leafKeyPath]] boolValue])
			[nodes addObjectsFromArray:[node descendantNodes]];
	}
	return [[nodes copy] autorelease];
}

- (NSTreeNode *)treeNodeAtIndexPath:(NSIndexPath *)indexPath; {
	return [[self arrangedObjects] descendantNodeAtIndexPath:indexPath];
}

- (NSArray *)treeNodesAtIndexPaths:(NSArray *)indexPaths; {
	NSMutableArray *nodes = [NSMutableArray array];
	for (NSIndexPath *path in indexPaths)
		[nodes addObject:[self treeNodeAtIndexPath:path]];
	return [[nodes copy] autorelease];
}

- (NSIndexPath *)indexPathForRepresentedObject:(id)object; {
	NSTreeNode *node = [self treeNodeForRepresentedObject:object];
	return (!node)?nil:[node indexPath];
}

- (NSArray *)indexPathsForRepresentedObjects:(NSArray *)objects; {
	NSMutableArray *indexPaths = [NSMutableArray array];
	for (id object in objects)
		[indexPaths addObject:[self indexPathForRepresentedObject:object]];
	return [[indexPaths copy] autorelease];
}

- (void)selectRepresentedObject:(id)object; {
	[self setSelectionIndexPath:[self indexPathForRepresentedObject:object]];
}

- (void)selectRepresentedObjects:(NSArray *)objects; {
	NSArray *nodes = [self flattenedNodes];
	NSMutableArray *paths = [NSMutableArray array];
	
	for (id object in objects) {
		for (NSTreeNode *node in nodes) {
			if ([[node representedObject] isEqualTo:object]) {
				[paths addObject:[node indexPath]];
				break;
			}
		}
	}		
	
	[self setSelectionIndexPaths:paths];
}

- (void)selectTreeNode:(NSTreeNode *)node; {
	[self setSelectionIndexPath:[node indexPath]];
}

- (void)selectTreeNodes:(NSArray *)nodes; {
	[self setSelectionIndexPaths:[nodes valueForKey:@"indexPath"]];
}

- (void)removeSelectedObjects; {
	[self removeObjectsAtArrangedObjectIndexPaths:[self selectionIndexPaths]];
}

- (void)removeTreeNode:(NSTreeNode *)node; {
	[self removeObjectAtArrangedObjectIndexPath:[node indexPath]];
}

- (void)removeTreeNodes:(NSArray *)nodes; {
	[self removeObjectsAtArrangedObjectIndexPaths:[nodes valueForKey:@"indexPath"]];
}

- (void)removeRepresentedObject:(id)object; {
	[self removeTreeNode:[self treeNodeForRepresentedObject:object]];
}

- (void)removeRepresentedObjects:(NSArray *)objects; {
	NSArray *nodes = [self flattenedNodes];
	NSMutableArray *paths = [NSMutableArray array];
	
	for (id object in objects) {
		for (NSTreeNode *node in nodes) {
			if ([[node representedObject] isEqualTo:object]) {
				[paths addObject:[node indexPath]];
				break;
			}
		}
	}
	
	[self removeObjectsAtArrangedObjectIndexPaths:paths];
}

- (void)updateChildNodeIndexesForTreeNode:(NSTreeNode *)item inProject:(WCProject *)project; {
	NSMutableIndexSet *indexes = [NSMutableIndexSet indexSet];
	NSMutableArray *objectsToMove = [NSMutableArray array];
	
	if (!item) {
		NSMutableArray *realNodes = [project projectGroups];
		NSArray *proxyNodes = [[self arrangedObjects] childNodes];
		
		for (NSTreeNode *node in proxyNodes) {
			NSUInteger realIndex = [realNodes indexOfObject:[node representedObject]];
			NSUInteger index = [proxyNodes indexOfObject:node];
			
			if (realIndex != index) {
				[objectsToMove addObject:[node representedObject]];
				[indexes addIndex:realIndex];
			}
		}
		
		[realNodes moveObjects:objectsToMove toIndexes:indexes];
	}
	else {
		for (NSTreeNode *node in [item childNodes]) {
			NSUInteger realIndex = [[[item representedObject] childNodes] indexOfObject:[node representedObject]];
			NSUInteger index = [[item childNodes] indexOfObject:node];
			
			if (realIndex != index) {
				[objectsToMove addObject:[node representedObject]];
				[indexes addIndex:realIndex];
			}
		}
		
		[[[item representedObject] mutableChildNodes] moveObjects:objectsToMove toIndexes:indexes];
	}
}
@end
